Open Domain Generalization with Domain-Augmented Meta-Learning

基于领域增强元学习的开放领域泛化

Yang Shu*, Zhangjie Cao*, Chenyu Wang, Jianmin Wang, Mingsheng Long ( )

School of Software, BNRist, Tsinghua University, China

{shu-y18,caozj14,cy-wang18}@mails.tsinghua.edu.cn, {jimwang,mingsheng}@tsinghua.edu.cn

Leveraging datasets available to learn a model with high generalization ability to unseen domains is important for computer vision, especially when the unseen domain's annotated data are unavailable. We study a novel and practical problem of Open Domain Generalization (OpenDG), which learns from different source domains to achieve high performance on an unknown target domain, where the distributions and label sets of each individual source domain and the target domain can be different. The problem can be generally applied to diverse source domains and widely applicable to real-world applications. We propose a Domain-Augmented Meta-Learning framework to learn open-domain generalizable representations. We augment domains on both feature-level by a new Dirichlet mixup and label-level by distilled soft-labeling, which complements each domain with missing classes and other domain knowledge. We conduct meta-learning over domains by designing new meta-learning tasks and losses to preserve domain unique knowledge and generalize knowledge across domains simultaneously. Experiment results on various multi-domain datasets demonstrate that the proposed Domain-Augmented Meta-Learning (DAML) outperforms prior methods for unseen domain recognition.

Deep convolutional neural networks have achieved state-of-the-art performance on wide ranges of computer vision applications with access to large-scale labeled data $[23,20$ , 39, 19]. However, for a target domain of interest, collecting enough training data is prohibitive. A practical solution is to generalize the model learned on the existing data to the unseen domain. Since the existing source datasets for training may be from different resources, they may fall into different domains and hold different label sets, e.g., ImageNet [8] and DomainNet [36]. Besides, the target domain is totally unknown, and may also have a distribution shift and a different label set from the source domains. We call the valuable and challenging problem as Open Domain Generalization

(OpenDG), where we need to learn generalizable representation from disparate source domains that generalizes well to any unseen target domain, as illustrated in Figure 1.

Figure 1. Open Domain Generalization (OpenDG). Different source domains hold disparate label sets. The goal is to learn generalizable representations from these source domains to help classify the known classes and detect open classes in the unseen target domain.

There are two key challenges for open domain generalization. (1) Distinct source domains and the unseen target domain are drawn from different distributions with a large distribution shift. (2) The different label sets of distinct source domains cause some classes to exist in many more domains than other classes. The data of minor classes existing in few domains are lacking in diversity. This makes the problem extremely difficult for existing methods [25, 29].

To address the first challenge, previous works minimize the distribution distance between domains by adversarial learning $\left[34,29\right]$ ,which successfully closes the domain gap when all source domains share the same label set. However, according to the second challenge, the different label sets between domains cause these distribution alignment methods to suffer from severe mismatch of classes. For the second challenge, a straightforward way is to manually sample data of minor classes existing in few domains, but the diversity in domains of the class is still limited. The generalization on the minor class is still inferior to other classes.

To generalize from arbitrary source domains to an unseen target domain, we propose a Domain-Augmented Meta-

*Equal contribution.

Table 1. Comparison of the proposed generalization setting with the previous settings related to cross-domain learning. The columns list assumptions made by the problem settings. Note that more " $x$ " means the method needs less assumption and thus is more widely-applicable. We can observe that the proposed open domain generalization problem requires no assumptions on the label set, no target data, and no post-training on target data,which is the most general problem setting. S means source while T means target. Note that "Same between S&T Domains" means the union of all source domain label sets equals the target label set, i.e., whether there are open classes.

Learning (DAML) framework. To close the domain gap between disparate source domains, we avoid distribution matching but learn generalizable representations across domains by meta-learning. To overcome the disparate label sets of open domain generalization, we propose two domain augmentation methods at both feature-level and label-level. At feature-level, we design a novel Dirichlet mixup (Dir-mixup) to compensate for the missing labels. At label-level, we utilize the soft-labeling distilled from other domains' networks to transfer the knowledge of other domains to the current network. DAML learns a representation that embeds the knowledge of all source domains and is highly generalizable to the unseen target domain. We use the ensemble of all source domain network outputs as the final prediction, which naturally calibrates the predictive uncertainty. In summary:

In this section, we briefly discuss works related to ours, including domain adaptation, domain generalization, and data augmentation methods. We compare our problem setting with the problem settings of previous works in Table 1.

Domain Adaptation aims to adapt the model from the source domain to the target domain, which typically mitigates the domain gap by minimizing the distribution distance $\left[14,32\right]$ . However,the classic domain adaptation requires the same label set between source and target domains. Recent works try to extend domain adaptation to varied source and target label sets $\left[2,35,41,51\right]$ ,but the solution relies on the target unlabeled data, which is not available in the open domain generalization setting.

Multi-source domain adaptation is more related to our work with more than one source domain. Most of the works assume that all the source domains share the same label set $\left[55,36\right]$ ,which can be easily violated in practice since source domains may be drawn from different resources. DCN [50] moves a step forward to remove the constraint on the source label sets but still requires the union of source label sets to be the same as the target label set. We instead require no label set constraint and no target data for training.

Domain Generalization aims to learn a generalizable model with only source data to achieve high performance in an unseen target domain $\left[22,34\right]$ ,which typically learns domain-invariant features across source domains [34, 16, $15,28,4,38,5]$ . When the different source domains hold different label sets, such learning causes mismatch of classes. CIDDG [29] can avoid the mismatching but still requires all the source and target domains to share the same label sets, or otherwise the low domain diversity of some classes makes it hard to learn domain-invariant features.

Meta-learning instead has the potential to learn from highly diverse domains. However, current meta-learning-based domain generalization methods still fail to consider different label sets of distinct source domains and the open classes in the target domain [25, 1, 10, 27]. Heterogeneous domain generalization $\left[30,49\right]$ has a similar goal of learning generalizable representations, which targets a more powerful pre-trained model by learning from heterogeneous source domains of different label sets. However, it requires additional target labeled data to induce a category model, which cannot fit into the proposed open domain generalization problem.

Augmentation The statistical learning theory [45] suggests that the generalization of the learning model can be characterized by the model capacity and the diversity of training data. So data augmentation can improve generalization by increasing the diversity of training data. Basic augmentations including affine transformation, random cropping, and horizontal flipping are widely-used in image classification [6, 42, 24]. Recently, more advanced augmentations are proposed. Mixup $\left[54,44,18\right]$ combines two samples linearly. Cutout [9] removes contiguous sections of input images. Cutmix [52] combines cutout and mixup by filling the Cutout part with sections of other image patches.

Algorithm 1 Training process of Domain-Augmented Meta-Learning (DAML)

Input: Source datasets ${\mathcal{D}}_1,{\mathcal{D}}_2,\cdots ,{\mathcal{D}}_S$ ,learning rates $\eta$ and $\beta$ ,Dir-mixup hyper-parameters ${\alpha }_{max}$ and ${\alpha }_{min}$

Initialize ${{\theta }_s|}_{s=1}^S$

while Not Converged do

Sample a batch of data ${\mathcal{B}}^{\mathrm{tr}}=\left\{\left({\mathbf{x}}_1,{\mathbf{y}}_1\right),\left({\mathbf{x}}_2,{\mathbf{y}}_2\right),\cdots ,\left({\mathbf{x}}_S,{\mathbf{y}}_S\right)\right\}$ from all source domains ${\mathcal{D}}_1,{\mathcal{D}}_2,\cdots ,{\mathcal{D}}_S$ .

for $s=1,\dots ,S$ do $⊳$ Meta-training starts

${\alpha }_s^{\mathrm{tr}}\leftarrow \left\{{\alpha }_{max},{\alpha }_{min},s\right\}$ $⊳$ Dir-mixup parameter for meta-training

${\mathcal{B}}_s^{\mathrm{D}\text{-mix }}=\left\{\left({\mathbf{z}}_s^{\mathrm{D}\text{-mix }},{\mathbf{y}}_s^{\mathrm{D}\text{-mix }}\right)\right\}\leftarrow \mathrm{Dir-mixup}\left(\left\{{\alpha }_s^{\mathrm{tr}},{\mathcal{B}}^{\mathrm{tr}}\right\}\right)⊳$ Obtain Dir-mixup according to Eqn. (3)

${\mathcal{B}}_s^{\text{distill }}=\left\{\left({\mathbf{x}}_s,{\mathbf{y}}_s^{\text{distill }}\right)\right\}\leftarrow \left\{{G_j|}_{j\ne s},{F_j|}_{j\ne s},{\mathcal{B}}^{\mathrm{tr}}\right\}$ $⊳$ Obtain distilled soft-label according to Eqn. (4)

${\mathcal{L}}_s^{\mathrm{tr}}\leftarrow \left\{G_s\left(F_s\left({\mathbf{x}}_s\right)\right),{\mathbf{y}}_s,G_s\left({\mathbf{z}}_s^{\text{D-mix }}\right),{\mathbf{y}}_s^{\text{D-mix }},{\mathbf{y}}_s^{\text{distill }}\right\}$ using data in ${\mathcal{B}}^{\mathrm{tr}},{\mathcal{B}}_s^{\text{D-mix }}$ ,and ${\mathcal{B}}_s^{\text{distill }}⊳$ According to Eqn (1)

${\theta }_{F_s^{\prime },G_s^{\prime }}={\theta }_{F_s,G_s}-\eta {\mathrm{\nabla }}_{\theta }{\mathcal{L}}_s^{\mathrm{tr}}$

Sample another batch of data ${\mathcal{B}}^{\text{obj }}=\left\{\left({\mathbf{x}}_1,{\mathbf{y}}_1\right),\left({\mathbf{x}}_2,{\mathbf{y}}_2\right),\cdots ,\left({\mathbf{x}}_S,{\mathbf{y}}_S\right)\right\}$ from all source domains ${\mathcal{D}}_1,{\mathcal{D}}_2,\cdots ,{\mathcal{D}}_S$ .

for $s=1,\dots ,S$ do $⊳$ Meta-objective starts

${\alpha }_s^{\text{obj }}\leftarrow \left\{{\alpha }_{min},{\alpha }_{max},s\right\}$ $⊳$ Dir-mixup parameter for meta-objective

${\mathcal{B}}_s^{\mathrm{D}\text{-mix }{}^{\prime }}=\left\{\left({\mathbf{z}}_s^{\mathrm{D}\text{-mix }{}^{\prime }},{\mathbf{y}}_s^{\mathrm{D}\text{-mix }{}^{\prime }}\right)\right\}\leftarrow \mathrm{Dir}\mathrm{-mixup}\left(\left\{{\alpha }_s^{\text{obj }},{\mathcal{B}}^{\text{obj }}\right\}\right)$ $⊳$ Obtain Dir-mixup according to Eqn. (3)

${\mathcal{L}}_s^{\text{obj }}\leftarrow \left\{G_s^{\prime }\left(F_s^{\prime }\left({\mathbf{x}}_j\right)\right){|{}_{j\ne s},{\mathbf{y}}_j|}_{j\ne s},G_s^{\prime }\left({\mathbf{z}}_s^{\mathrm{D}-{\mathrm{mix}}^{\prime }}\right),{\mathbf{y}}_s^{\mathrm{D}-{\mathrm{mix}}^{\prime }}\right\}$ using data in ${\mathcal{B}}^{\text{obj }}$ and ${\mathcal{B}}_s^{\mathrm{D}-{\mathrm{mix}}^{\prime }}⊳$ According to Eqn (2)

${\theta }_{F_s,G_s}\leftarrow {\theta }_{F_s,G_s}-\beta {\mathrm{\nabla }}_{\theta }\left({\mathcal{L}}_s^{\mathrm{tr}}+{\mathcal{L}}_s^{\mathrm{obj}}\right)$ $⊳$ Update parameters with meta-learning

return ${{\theta }_s|}_{s=1}^S$

Augmentation-based generalization methods promote the generalization ability by augmenting source data, where adversarial data augmentation [47], gradient-based perturbations [43], self-supervised learning signals [3], and Cut-Mix [33] are used as the augmentation method. Note that these augmentation methods target general situations for generalization across domains but are not designed specially for open domains with disparate label sets.

Different from all previous works, this paper studies open domain generalization, a practical but challenging problem. We develop the DAML framework to conduct meta-learning over augmented source domains. We design a novel Dir-mixup to mix samples from multiple domains instead of mixing two arbitrary samples in classic mixup. Dir-mixup bridges all the source domains and compensates each domain with missing classes from other domains, which naturally fits the disparate source label sets. We further propose a new distilled soft-labeling to transfer knowledge across domains.

In this section, we first introduce the open domain generalization (OpenDG) problem. Then we introduce the Domain-Augmented Meta-Learning (DAML) and describe the step-by-step algorithm and the optimization of the framework, which consists of the proposed domain augmentation and the meta-learning on the augmented domains.

In open domain generalization (OpenDG), we have multiple source domains ${\mathcal{D}}_1,{\mathcal{D}}_2,\cdots ,{\mathcal{D}}_S$ available for training, where each source domain $s$ consists of data-label pairs ${\mathcal{D}}_s=\left\{\left({\mathbf{x}}_s,{\mathbf{y}}_s\right)\right\}.{\mathbf{y}}_s$ denotes the one-hot label of ${\mathbf{x}}_s$ . Note that although we train the model with mini-batches in practice, here we omit the batch size of each domain to simplify the notations. We use $\mathcal{C}$ to denote the union of all the source label sets. In open domain generalization, we have no constraint on the label sets of different domains. We aim to learn open-domain generalizable representation from all the source domains, which can generalize well to an unseen target domain ${\mathcal{D}}_t$ . Specifically,the target domain,only used for evaluation,consists of fully unlabeled data ${\mathcal{D}}_t=\left\{{\mathbf{x}}_t\right\}$ and its label set ${\mathcal{C}}_t$ may contain classes existing in any source label set or unknown classes not existing in the union of source label sets $\mathcal{C}$ . The goal is to classify at inference each target sample with the correct class if it belongs to the source label set $\mathcal{C}$ ,or label it as "unknown". Note that no target data,even unlabeled, are available for training, which differs OpenDG from domain adaptation [51] or domain generalization [49].

Figure 2. The architecture of the proposed DAML framework. We show the computation graph for source domain 1 as an example, and the other source domains are computed similarly. In the meta-training (up part, left to right), each source domain is augmented by Dir-mixup (red) and distilled soft-labeling (blue) to compute the ${\mathcal{L}}_1^{\mathrm{tr}}$ to update the model parameters to $F_1^{\prime }$ and $G_1^{\prime }$ . In the meta-objective (down part, right to left),each source domain is augmented by Dir-mixup (red) to compute the ${\mathcal{L}}_1^{\text{obj }}$ to finally update the model parameters.

We propose DAML to address open domain generalization problems to mitigate the disparate label sets and distribution shifts among the diverse source domains. As shown in Algorithm 1, the idea is to learn generalizable representations by meta-learning over augmented domains.

Augmented Domains As demonstrated in [53, 17], increasing the diversity of the dataset can substantially improve the generalization of the representations. Motivated by this idea, we augment each domain to expand the diversity of the datasets. We observe that different domains have different distributions and hold different label sets, which means that each domain contains distinct knowledge but lacks domain knowledge and class knowledge of other domains. Based on the observation, we design domain augmentation to address open domain generalization. Our insight is to conduct both feature-level and label-level augmentation. For feature-level augmentation, we propose a novel Dirichlet Mixup (Dir-mixup) method, which augments each domain by the mixup with multiple domains. For label-level augmentation, we propose to augment each domain by distilling soft-labels from models of other domains. The proposed domain augmentation increases the diversity of the data and compensates each domain with missing knowledge of features and classes. The details of the proposed domain augmentation are introduced in Section 3.3.

Meta-Learning We design the learning framework to learn generalizable representations, which simultaneously preserves the unique information of each domain and aggregates the knowledge of all the domains. Thus, instead of employing a shared network for all source domains, which only embeds domain common knowledge, we build one individual classification network composed of a feature extractor $F_s$ and a classifier $G_s$ for each source domain $s$ . Then we need to learn a generalizable representation aggregating the information of all the source domains. We conduct meta-learning over all the networks since meta-learning is demonstrated to be able to learn a generalizable representation from highly disparate domains. In each iteration of the parameter update, we first draw a batch of samples from each domain and compute the corresponding Dir-mixup samples and distilled soft-labels (Line 5-7 in Algorithm 1). Unlike standard meta-learning loss applied only on the raw data [12], with the augmented domains, we design a new meta-training loss as the classification loss on the original data, the domain-augmented data by Dir-mixup, and soft-labels distilled from other domain networks. For each domain $s$ ,let ${\mathbf{z}}_s=F_s\left({\mathbf{x}}_s\right)$ be the feature of ${\mathbf{x}}_s$ ,we define the meta-training loss as

(1)

The superscript(k)means the probability of the $k$ -th class. ${\mathcal{D}}_s^{\text{D-mix }}$ and ${\mathcal{D}}_s^{\text{distill }}$ are the augmented domains of Dir-mixup samples and distilled soft-label samples for meta-training on domain $s$ . We compute one step of gradient update for each source network with respect to the meta-training loss: ${\theta }_{G_s^{\prime },F_s^{\prime }}={\theta }_{G_s,F_s}-\eta {\mathrm{\nabla }}_{\theta }{\mathcal{L}}_s^{\mathrm{tr}}$ (Line 9 in Algorithm 1),where $\eta$ is the step size. The design idea of meta-objective is to guide the gradient update from the meta-training loss to the desired goal. Classic meta-learning employs the losses over all sampled tasks as the meta-objective [12]. But our goal is to improve the generalization ability of the model, so different from classic meta-objective, we design the meta-objective as the classification loss on the original data and Dir-mixup data in other domains with the updated network $G_s^{\prime },F_s^{\prime }$ ,which can propagate the knowledge of other domains to domain $s$ and promote the knowledge transfer and generalization across domains. The meta-objective is defined as

(2)

${\mathcal{D}}_s^{\mathrm{D}\text{-mix }}{}^{\prime }$ is the augmented domain of Dir-mixup samples for domain $s$ in meta-objective. The minimization of the meta-objective finds a gradient descent update that updates the network to classify data in other domains with high accuracy, which encourages the network to learn a generalizable representation performing well across all domains. We finally update the network parameters in one iteration by ${\theta }_s\leftarrow {\theta }_s-\beta {\mathrm{\nabla }}_{\theta }\left({\mathcal{L}}_s^{\text{tr }}+{\mathcal{L}}_s^{\text{obj }}\right)$ ,where $\beta$ is the learning rate.

The meta-learning framework can learn a generalizable representation aggregating information from all source domains, where the generalization power highly relies on the diversity of each source domain. To this end, we design two multiple source domain augmentation approaches: the feature-level augmentation, Dir-mixup, and the label-level augmentation, distilled augmentation. The augmentations compensate for the missing class information in each source domain and further increase domain diversity.

Dir-mixup Mixup [54] generates a new data-label by the weighted sum of the feature and one-hot label of existing samples, where the weights are sampled from a pre-defined distribution. We augment the $s$ -th source domain by mixup of data in the $s$ -th domain with data in other domains. Since these data may belong to the missing classes of the $s$ -th source domain, mixup augmentation would compensate for the missing classes. Also, mixup produces inter-domain data, which further increases the diversity of data in each domain.

However, the original mixup is defined to mix two samples. When applied to open domain generalization with multiple source domains, mixup samples are only generated from pairs of domains, which, as shown in Figure 3, only generates samples between two domains (the lines between vertex) but lacks samples mixing multiple domains (the whole area). Also, to obtain all domain combinations, such pairwise mixup needs $O$ (#domains $\times$ #domains) mixup samples. Therefore, to mix multiple domains, we need to sample the weight from a multi-variate distribution instead of the beta distribution used in the original mixup. We select Dirichlet distribution since it has similar properties to the beta distribution and is a multi-variate distribution. We then design a new Dir-mixup to mix samples (one for each domain) with a designed weight $\lambda$ sampled from a Dirichlet distribution parameterized by a parameter $\alpha$ . We perform mixup at feature-level. Let ${\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_S$ be the features of different domain data extracted by the network, the Dir-mixup augmented data $\left({\mathbf{z}}^{\mathrm{D}\text{-mix }},{\mathbf{y}}^{\mathrm{D}\text{-mix }}\right)$ can be calculated as:

Figure 3. Comparison between Dir-mixup and classic mixup. Classic mixup only mixes two samples, so mixup samples only exist on the edge of the triangle while Dir-mixup mixes samples of multiple domains covering the whole triangle area, meaning Dir-mixup introduce mixup samples with more information and higher diversity.

Compared with recent work using mixup for domain generalization $\left[33,49\right]$ ,Dir-mixup is more efficient and effective. The parameter $\alpha$ adjusts the distribution to generate different augmentations, better serving the meta-learning process. Consider constructing Dir-mixup for each model $s$ . In the meta-training, we want to keep more information and focus more on domain $s$ during mixup,so we set ${\alpha }^{\left(s\right)}$ larger than other components in $\alpha$ ,which assigns a larger weight ${\lambda }^{\left(s\right)}$ to ${\mathbf{z}}_s$ statistically. In the meta objective,the goal is to transfer knowledge from other domains and improve the cross-domain generalization, which would be enhanced by mixup results with larger domain discrepancy. So we set ${\alpha }^{\left(s\right)}$ smaller than other components in $\alpha$ ,which induces smaller ${\lambda }^{\left(s\right)}$ statistically. We employ two hyper-parameters ${\alpha }_{max}$ and ${\alpha }_{min}$ to realize this idea. For the meta-training of model $s$ ,we set ${\alpha }_s^{\mathrm{tr}}$ to be a length $S$ vector with all entries as ${\alpha }_{min}$ but the $s$ -th entry as ${\alpha }_{max}$ . We generate mixup data with this ${\alpha }_s^{\mathrm{tr}}$ to form the Dir-mixup augmentation set in the meta-training of model $s$ ,as ${\mathcal{D}}_s^{\mathrm{D}\text{-mix }}$ in Equation 1. For the meta-objective,we set ${\alpha }_s^{\text{obj }}$ to be a length $S$ vector with all entries as ${\alpha }_{max}$ but the $s$ -th entry as ${\alpha }_{min}$ . And the data generated from this ${\alpha }_s^{\text{obj }}$ form the Dir-mixup augmentation set for model $s$ ,which is the ${\mathcal{D}}_s^{\mathrm{D}\text{-mix’ }}$ in Equation 2.